package com.LeeCode;

import java.util.*;
import java.util.stream.Collectors;

/**
 * 前K个高频元素
 */

public class Code347 {
    class Node {
        int data;
        int count;

        public Node(int data, int count) {
            this.data = data;
            this.count = count;
        }
    }

    public static void main(String[] args) {
        int[] nums = {1, 1, 1, 2, 2, 3};
        System.out.println(Arrays.toString(new Code347().topKFrequent(nums, 2)));
    }

    public int[] topKFrequent(int[] nums, int k) {
        HashMap<Integer, Node> map = new HashMap<>();
        int[] res = new int[k];

        for (int i = 0; i < nums.length; i++) {
            if (!map.containsKey(nums[i])) {
                map.put(nums[i], new Node(nums[i], 1));
            } else {
                map.get(nums[i]).count++;
            }
        }
        ArrayList<Node> list = new ArrayList<>(map.values());
        int n = list.size();
        for (int i = n / 2 - 1; i >= 0; i--) {
            heapfy(list, n, i);
        }

        for (int i = 0; i < k; i++) {
            Collections.swap(list, 0, n - 1 - i);
            heapfy(list, n - 1 - i, 0);
            res[i] = list.get(n - 1 - i).data;
        }
        return res;
    }

    public void heapfy(ArrayList<Node> list, int heapSize, int rootIndex) {
        int largest = rootIndex;
        int leftChild = 2 * rootIndex + 1;
        int rightChild = 2 * rootIndex + 2;

        if (leftChild < heapSize && list.get(leftChild).count > list.get(largest).count) {
            largest = leftChild;
        }

        if (rightChild < heapSize && list.get(rightChild).count > list.get(largest).count) {
            largest = rightChild;
        }

        if (largest != rootIndex) {
            Collections.swap(list, rootIndex, largest);
            heapfy(list, heapSize, largest);
        }
    }
}
